// Ceres Solver - A fast non-linear least squares minimizer
// Copyright 2023 Google Inc. All rights reserved.
// http://ceres-solver.org/
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice,
//   this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
// * Neither the name of Google Inc. nor the names of its contributors may be
//   used to endorse or promote products derived from this software without
//   specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Author: sameeragarwal@google.com (Sameer Agarwal)
//
// Generic loop for line search based optimization algorithms.
//
// This is primarily inspired by the minFunc packaged written by Mark
// Schmidt.
//
// http://www.di.ens.fr/~mschmidt/Software/minFunc.html
//
// For details on the theory and implementation see "Numerical
// Optimization" by Nocedal & Wright.

#include "ceres/line_search_minimizer.h"

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>

#include "Eigen/Dense"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/strings/str_format.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "ceres/array_utils.h"
#include "ceres/evaluator.h"
#include "ceres/internal/eigen.h"
#include "ceres/internal/export.h"
#include "ceres/line_search.h"
#include "ceres/line_search_direction.h"
#include "ceres/types.h"

namespace ceres::internal {
namespace {

bool EvaluateGradientNorms(Evaluator* evaluator,
                           const Vector& x,
                           LineSearchMinimizer::State* state,
                           std::string* message) {
  Vector negative_gradient = -state->gradient;
  Vector projected_gradient_step(x.size());
  if (!evaluator->Plus(
          x.data(), negative_gradient.data(), projected_gradient_step.data())) {
    *message = "projected_gradient_step = Plus(x, -gradient) failed.";
    return false;
  }

  state->gradient_squared_norm = (x - projected_gradient_step).squaredNorm();
  state->gradient_max_norm =
      (x - projected_gradient_step).lpNorm<Eigen::Infinity>();
  return true;
}

}  // namespace

void LineSearchMinimizer::Minimize(const Minimizer::Options& options,
                                   double* parameters,
                                   Solver::Summary* summary) {
  const bool is_not_silent = !options.is_silent;
  const absl::Time start_time = absl::Now();
  absl::Time iteration_start_time = start_time;

  CHECK(options.evaluator != nullptr);
  Evaluator* evaluator = options.evaluator.get();
  const int num_parameters = evaluator->NumParameters();
  const int num_effective_parameters = evaluator->NumEffectiveParameters();

  summary->termination_type = NO_CONVERGENCE;
  summary->num_successful_steps = 0;
  summary->num_unsuccessful_steps = 0;

  VectorRef x(parameters, num_parameters);

  State current_state(num_parameters, num_effective_parameters);
  State previous_state(num_parameters, num_effective_parameters);

  IterationSummary iteration_summary;
  iteration_summary.iteration = 0;
  iteration_summary.step_is_valid = false;
  iteration_summary.step_is_successful = false;
  iteration_summary.cost_change = 0.0;
  iteration_summary.gradient_max_norm = 0.0;
  iteration_summary.gradient_norm = 0.0;
  iteration_summary.step_norm = 0.0;
  iteration_summary.linear_solver_iterations = 0;
  iteration_summary.step_solver_time_in_seconds = 0;

  // Do initial cost and gradient evaluation.
  if (!evaluator->Evaluate(x.data(),
                           &(current_state.cost),
                           nullptr,
                           current_state.gradient.data(),
                           nullptr)) {
    summary->termination_type = FAILURE;
    summary->message = "Initial cost and jacobian evaluation failed.";
    if (is_not_silent) {
      LOG(WARNING) << "Terminating: " << summary->message;
    }
    return;
  }

  if (!EvaluateGradientNorms(evaluator, x, &current_state, &summary->message)) {
    summary->termination_type = FAILURE;
    summary->message =
        "Initial cost and jacobian evaluation failed. More details: " +
        summary->message;
    if (is_not_silent) {
      LOG(WARNING) << "Terminating: " << summary->message;
    }
    return;
  }

  summary->initial_cost = current_state.cost + summary->fixed_cost;
  iteration_summary.cost = current_state.cost + summary->fixed_cost;

  iteration_summary.gradient_norm = sqrt(current_state.gradient_squared_norm);
  iteration_summary.gradient_max_norm = current_state.gradient_max_norm;
  if (iteration_summary.gradient_max_norm <= options.gradient_tolerance) {
    summary->message = absl::StrFormat(
        "Gradient tolerance reached. Gradient max norm: %e <= %e",
        iteration_summary.gradient_max_norm,
        options.gradient_tolerance);
    summary->termination_type = CONVERGENCE;
    if (is_not_silent) {
      VLOG(1) << "Terminating: " << summary->message;
    }
    return;
  }

  absl::Time now = absl::Now();
  iteration_summary.iteration_time_in_seconds =
      absl::ToDoubleSeconds(now - iteration_start_time);
  iteration_summary.cumulative_time_in_seconds =
      absl::ToDoubleSeconds(now - start_time) +
      summary->preprocessor_time_in_seconds;
  summary->iterations.push_back(iteration_summary);

  LineSearchDirection::Options line_search_direction_options;
  line_search_direction_options.num_parameters = num_effective_parameters;
  line_search_direction_options.type = options.line_search_direction_type;
  line_search_direction_options.nonlinear_conjugate_gradient_type =
      options.nonlinear_conjugate_gradient_type;
  line_search_direction_options.max_lbfgs_rank = options.max_lbfgs_rank;
  line_search_direction_options.use_approximate_eigenvalue_bfgs_scaling =
      options.use_approximate_eigenvalue_bfgs_scaling;
  std::unique_ptr<LineSearchDirection> line_search_direction =
      LineSearchDirection::Create(line_search_direction_options);

  LineSearchFunction line_search_function(evaluator);

  LineSearch::Options line_search_options;
  line_search_options.interpolation_type =
      options.line_search_interpolation_type;
  line_search_options.min_step_size = options.min_line_search_step_size;
  line_search_options.sufficient_decrease =
      options.line_search_sufficient_function_decrease;
  line_search_options.max_step_contraction =
      options.max_line_search_step_contraction;
  line_search_options.min_step_contraction =
      options.min_line_search_step_contraction;
  line_search_options.max_num_iterations =
      options.max_num_line_search_step_size_iterations;
  line_search_options.sufficient_curvature_decrease =
      options.line_search_sufficient_curvature_decrease;
  line_search_options.max_step_expansion =
      options.max_line_search_step_expansion;
  line_search_options.is_silent = options.is_silent;
  line_search_options.function = &line_search_function;

  std::unique_ptr<LineSearch> line_search(LineSearch::Create(
      options.line_search_type, line_search_options, &summary->message));
  if (line_search.get() == nullptr) {
    summary->termination_type = FAILURE;
    if (is_not_silent) {
      LOG(ERROR) << "Terminating: " << summary->message;
    }
    return;
  }

  LineSearch::Summary line_search_summary;
  int num_line_search_direction_restarts = 0;

  while (true) {
    if (!RunCallbacks(options, iteration_summary, summary)) {
      break;
    }

    iteration_start_time = absl::Now();
    if (iteration_summary.iteration >= options.max_num_iterations) {
      summary->message = "Maximum number of iterations reached.";
      summary->termination_type = NO_CONVERGENCE;
      if (is_not_silent) {
        VLOG(1) << "Terminating: " << summary->message;
      }
      break;
    }

    const double total_solver_time =
        absl::ToDoubleSeconds(iteration_start_time - start_time) +
        summary->preprocessor_time_in_seconds;
    if (total_solver_time >= options.max_solver_time_in_seconds) {
      summary->message = "Maximum solver time reached.";
      summary->termination_type = NO_CONVERGENCE;
      if (is_not_silent) {
        VLOG(1) << "Terminating: " << summary->message;
      }
      break;
    }

    iteration_summary = IterationSummary();
    iteration_summary.iteration = summary->iterations.back().iteration + 1;
    iteration_summary.step_is_valid = false;
    iteration_summary.step_is_successful = false;

    bool line_search_status = true;
    if (iteration_summary.iteration == 1) {
      current_state.search_direction = -current_state.gradient;
    } else {
      line_search_status = line_search_direction->NextDirection(
          previous_state, current_state, &current_state.search_direction);
    }

    if (!line_search_status &&
        num_line_search_direction_restarts >=
            options.max_num_line_search_direction_restarts) {
      // Line search direction failed to generate a new direction, and we
      // have already reached our specified maximum number of restarts,
      // terminate optimization.
      summary->message = absl::StrFormat(
          "Line search direction failure: specified "
          "max_num_line_search_direction_restarts: %d reached.",
          options.max_num_line_search_direction_restarts);
      summary->termination_type = FAILURE;
      if (is_not_silent) {
        LOG(WARNING) << "Terminating: " << summary->message;
      }
      break;
    } else if (!line_search_status) {
      // Restart line search direction with gradient descent on first iteration
      // as we have not yet reached our maximum number of restarts.
      CHECK_LT(num_line_search_direction_restarts,
               options.max_num_line_search_direction_restarts);

      ++num_line_search_direction_restarts;
      if (is_not_silent) {
        LOG(WARNING) << "Line search direction algorithm: "
                     << LineSearchDirectionTypeToString(
                            options.line_search_direction_type)
                     << ", failed to produce a valid new direction at "
                     << "iteration: " << iteration_summary.iteration
                     << ". Restarting, number of restarts: "
                     << num_line_search_direction_restarts << " / "
                     << options.max_num_line_search_direction_restarts
                     << " [max].";
      }
      line_search_direction =
          LineSearchDirection::Create(line_search_direction_options);
      current_state.search_direction = -current_state.gradient;
    }

    line_search_function.Init(x, current_state.search_direction);
    current_state.directional_derivative =
        current_state.gradient.dot(current_state.search_direction);

    // TODO(sameeragarwal): Refactor this into its own object and add
    // explanations for the various choices.
    //
    // Note that we use !line_search_status to ensure that we treat cases when
    // we restarted the line search direction equivalently to the first
    // iteration.
    const double initial_step_size =
        (iteration_summary.iteration == 1 || !line_search_status)
            ? std::min(1.0, 1.0 / current_state.gradient_max_norm)
            : std::min(1.0,
                       2.0 * (current_state.cost - previous_state.cost) /
                           current_state.directional_derivative);
    // By definition, we should only ever go forwards along the specified search
    // direction in a line search, most likely cause for this being violated
    // would be a numerical failure in the line search direction calculation.
    if (initial_step_size < 0.0) {
      summary->message = absl::StrFormat(
          "Numerical failure in line search, initial_step_size is "
          "negative: %.5e, directional_derivative: %.5e, "
          "(current_cost - previous_cost): %.5e",
          initial_step_size,
          current_state.directional_derivative,
          (current_state.cost - previous_state.cost));
      summary->termination_type = FAILURE;
      if (is_not_silent) {
        LOG(WARNING) << "Terminating: " << summary->message;
      }
      break;
    }

    line_search->Search(initial_step_size,
                        current_state.cost,
                        current_state.directional_derivative,
                        &line_search_summary);
    if (!line_search_summary.success) {
      summary->message = absl::StrFormat(
          "Numerical failure in line search, failed to find "
          "a valid step size, (did not run out of iterations) "
          "using initial_step_size: %.5e, initial_cost: %.5e, "
          "initial_gradient: %.5e.",
          initial_step_size,
          current_state.cost,
          current_state.directional_derivative);
      if (is_not_silent) {
        LOG(WARNING) << "Terminating: " << summary->message;
      }
      summary->termination_type = FAILURE;
      break;
    }

    const FunctionSample& optimal_point = line_search_summary.optimal_point;
    CHECK(optimal_point.vector_x_is_valid)
        << "Congratulations, you found a bug in Ceres. Please report it.";
    current_state.step_size = optimal_point.x;
    previous_state = current_state;
    iteration_summary.step_solver_time_in_seconds =
        absl::ToDoubleSeconds(absl::Now() - iteration_start_time);

    if (optimal_point.vector_gradient_is_valid) {
      current_state.cost = optimal_point.value;
      current_state.gradient = optimal_point.vector_gradient;
    } else {
      Evaluator::EvaluateOptions evaluate_options;
      evaluate_options.new_evaluation_point = false;
      if (!evaluator->Evaluate(evaluate_options,
                               optimal_point.vector_x.data(),
                               &(current_state.cost),
                               nullptr,
                               current_state.gradient.data(),
                               nullptr)) {
        summary->termination_type = FAILURE;
        summary->message = "Cost and jacobian evaluation failed.";
        if (is_not_silent) {
          LOG(WARNING) << "Terminating: " << summary->message;
        }
        return;
      }
    }

    if (!EvaluateGradientNorms(evaluator,
                               optimal_point.vector_x,
                               &current_state,
                               &summary->message)) {
      summary->termination_type = FAILURE;
      summary->message =
          "Step failed to evaluate. This should not happen as the step was "
          "valid when it was selected by the line search. More details: " +
          summary->message;
      if (is_not_silent) {
        LOG(WARNING) << "Terminating: " << summary->message;
      }
      break;
    }

    // Compute the norm of the step in the ambient space.
    iteration_summary.step_norm = (optimal_point.vector_x - x).norm();
    const double x_norm = x.norm();
    x = optimal_point.vector_x;

    iteration_summary.gradient_max_norm = current_state.gradient_max_norm;
    iteration_summary.gradient_norm = sqrt(current_state.gradient_squared_norm);
    iteration_summary.cost_change = previous_state.cost - current_state.cost;
    iteration_summary.cost = current_state.cost + summary->fixed_cost;

    iteration_summary.step_is_valid = true;
    iteration_summary.step_is_successful = true;
    iteration_summary.step_size = current_state.step_size;
    iteration_summary.line_search_function_evaluations =
        line_search_summary.num_function_evaluations;
    iteration_summary.line_search_gradient_evaluations =
        line_search_summary.num_gradient_evaluations;
    iteration_summary.line_search_iterations =
        line_search_summary.num_iterations;
    const absl::Time now = absl::Now();
    iteration_summary.iteration_time_in_seconds =
        absl::ToDoubleSeconds(now - iteration_start_time);
    iteration_summary.cumulative_time_in_seconds =
        absl::ToDoubleSeconds(now - start_time) +
        summary->preprocessor_time_in_seconds;
    summary->iterations.push_back(iteration_summary);

    // Iterations inside the line search algorithm are considered
    // 'steps' in the broader context, to distinguish these inner
    // iterations from from the outer iterations of the line search
    // minimizer. The number of line search steps is the total number
    // of inner line search iterations (or steps) across the entire
    // minimization.
    summary->num_line_search_steps += line_search_summary.num_iterations;
    summary->line_search_cost_evaluation_time_in_seconds +=
        absl::ToDoubleSeconds(line_search_summary.cost_evaluation_time);
    summary->line_search_gradient_evaluation_time_in_seconds +=
        absl::ToDoubleSeconds(line_search_summary.gradient_evaluation_time);
    summary->line_search_polynomial_minimization_time_in_seconds +=
        absl::ToDoubleSeconds(line_search_summary.polynomial_minimization_time);
    summary->line_search_total_time_in_seconds +=
        absl::ToDoubleSeconds(line_search_summary.total_time);
    ++summary->num_successful_steps;

    const double step_size_tolerance =
        options.parameter_tolerance * (x_norm + options.parameter_tolerance);
    if (iteration_summary.step_norm <= step_size_tolerance) {
      summary->message = absl::StrFormat(
          "Parameter tolerance reached. "
          "Relative step_norm: %e <= %e.",
          (iteration_summary.step_norm /
           (x_norm + options.parameter_tolerance)),
          options.parameter_tolerance);
      summary->termination_type = CONVERGENCE;
      if (is_not_silent) {
        VLOG(1) << "Terminating: " << summary->message;
      }
      return;
    }

    if (iteration_summary.gradient_max_norm <= options.gradient_tolerance) {
      summary->message = absl::StrFormat(
          "Gradient tolerance reached. "
          "Gradient max norm: %e <= %e",
          iteration_summary.gradient_max_norm,
          options.gradient_tolerance);
      summary->termination_type = CONVERGENCE;
      if (is_not_silent) {
        VLOG(1) << "Terminating: " << summary->message;
      }
      break;
    }

    const double absolute_function_tolerance =
        options.function_tolerance * std::abs(previous_state.cost);
    if (std::abs(iteration_summary.cost_change) <=
        absolute_function_tolerance) {
      summary->message = absl::StrFormat(
          "Function tolerance reached. "
          "|cost_change|/cost: %e <= %e",
          std::abs(iteration_summary.cost_change) / previous_state.cost,
          options.function_tolerance);
      summary->termination_type = CONVERGENCE;
      if (is_not_silent) {
        VLOG(1) << "Terminating: " << summary->message;
      }
      break;
    }
  }
}

}  // namespace ceres::internal
